#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Prompt 统一加载器
支持加载和管理 AI Prompt，提供版本控制和工作流支持
"""
import hashlib
import yaml
from pathlib import Path
from typing import Dict, Any, Optional, List


class PromptLoader:
    """Prompt 加载器"""

    def __init__(self, prompts_dir: str = "prompts"):
        """
        初始化 Prompt 加载器

        Args:
            prompts_dir: Prompt 文件目录路径
        """
        self.prompts_dir = Path(prompts_dir)
        self.index_file = self.prompts_dir / "index.yaml"
        self.index = self._load_index()

    def _load_index(self) -> Dict[str, Any]:
        """加载 Prompt 索引文件"""
        if not self.index_file.exists():
            raise FileNotFoundError(f"Prompt index not found: {self.index_file}")

        with open(self.index_file, "r", encoding="utf-8") as f:
            return yaml.safe_load(f)

    def load_prompt(self, prompt_id: str) -> Dict[str, Any]:
        """
        加载指定 Prompt

        Args:
            prompt_id: Prompt 标识符（如 "ai_auto_self_diagnose"）

        Returns:
            {
                "id": "ai_auto_self_diagnose",
                "file": "PROMPT_005 问题诊断修复器.md",
                "role": "问题诊断专家",
                "content": "完整的 Prompt 内容",
                "parsed_config": {...},  # 解析的 YAML 配置
                "checksum": "md5 hash"
            }

        Raises:
            ValueError: 如果 prompt_id 不存在
            FileNotFoundError: 如果 Prompt 文件不存在
        """
        if prompt_id not in self.index.get("prompts", {}):
            available = list(self.index.get("prompts", {}).keys())
            raise ValueError(
                f"Prompt not found: {prompt_id}\n"
                f"Available prompts: {', '.join(available)}"
            )

        prompt_info = self.index["prompts"][prompt_id]
        prompt_file = self.prompts_dir / prompt_info["file"]

        if not prompt_file.exists():
            raise FileNotFoundError(f"Prompt file not found: {prompt_file}")

        # 读取 Prompt 文件内容
        content = prompt_file.read_text(encoding="utf-8")

        # 解析 Prompt 内容
        parsed_config = self._parse_prompt_content(content)

        # 计算内容校验和
        checksum = hashlib.md5(content.encode('utf-8')).hexdigest()

        return {
            "id": prompt_id,
            "file": prompt_info["file"],
            "version": prompt_info.get("version", "unknown"),
            "role": prompt_info.get("role", "AI Assistant"),
            "description": prompt_info.get("description", ""),
            "trigger": prompt_info.get("trigger", "manual"),
            "output_format": prompt_info.get("output_format", "text"),
            "content": content,
            "parsed_config": parsed_config,
            "checksum": checksum,
            "metadata": prompt_info
        }

    def _parse_prompt_content(self, content: str) -> Dict[str, Any]:
        """
        解析 Prompt 内容，提取 YAML Front Matter

        Args:
            content: Prompt 文件的完整内容

        Returns:
            解析后的配置字典
        """
        # 检查是否有 YAML Front Matter (```yaml ... ```)
        if content.strip().startswith("```yaml"):
            yaml_end = content.find("```", 7)
            if yaml_end != -1:
                yaml_content = content[7:yaml_end].strip()
                try:
                    return yaml.safe_load(yaml_content)
                except yaml.YAMLError as e:
                    print(f"[PromptLoader] 警告: YAML 解析失败: {e}")
                    return {}

        # 如果没有 YAML Front Matter，尝试提取 Markdown 标题作为结构
        return self._extract_markdown_structure(content)

    def _extract_markdown_structure(self, content: str) -> Dict[str, Any]:
        """
        从 Markdown 内容提取结构化信息

        Args:
            content: Markdown 内容

        Returns:
            提取的结构化信息
        """
        structure = {}
        lines = content.split('\n')

        for line in lines:
            # 提取 role（假设在开头）
            if line.startswith("role:"):
                structure["role"] = line.split(":", 1)[1].strip().strip('"')

        return structure

    def get_workflow(self, workflow_id: str) -> Dict[str, Any]:
        """
        获取工作流定义

        Args:
            workflow_id: 工作流标识符

        Returns:
            工作流定义字典

        Raises:
            ValueError: 如果工作流不存在
        """
        if "workflows" not in self.index:
            raise ValueError("No workflows defined in index")

        workflows = self.index["workflows"]
        if workflow_id not in workflows:
            available = list(workflows.keys())
            raise ValueError(
                f"Workflow not found: {workflow_id}\n"
                f"Available workflows: {', '.join(available)}"
            )

        return workflows[workflow_id]

    def list_prompts(self) -> List[Dict[str, str]]:
        """
        列出所有可用的 Prompt

        Returns:
            Prompt 信息列表
        """
        prompts = []
        for prompt_id, prompt_info in self.index.get("prompts", {}).items():
            prompts.append({
                "id": prompt_id,
                "file": prompt_info["file"],
                "role": prompt_info.get("role", "Unknown"),
                "description": prompt_info.get("description", ""),
                "trigger": prompt_info.get("trigger", "manual")
            })
        return prompts

    def list_workflows(self) -> List[Dict[str, str]]:
        """
        列出所有可用的工作流

        Returns:
            工作流信息列表
        """
        workflows = []
        for workflow_id, workflow_info in self.index.get("workflows", {}).items():
            workflows.append({
                "id": workflow_id,
                "name": workflow_info.get("name", workflow_id),
                "description": workflow_info.get("description", ""),
                "steps": len(workflow_info.get("steps", []))
            })
        return workflows

    def build_system_prompt(self, prompt_id: str) -> str:
        """
        构建用于 AI 的 system prompt

        Args:
            prompt_id: Prompt 标识符

        Returns:
            格式化的 system prompt 字符串
        """
        prompt_data = self.load_prompt(prompt_id)
        config = prompt_data["parsed_config"]

        # 构建结构化 system prompt
        role = config.get("role", prompt_data["role"])

        system_prompt = f"""You are a {role}.

Your responsibilities and expertise:
"""

        # 添加专业能力
        if "diagnostic_expertise" in config:
            system_prompt += "\nDiagnostic Expertise:\n"
            for key, value in config["diagnostic_expertise"].items():
                system_prompt += f"- {key}: {value}\n"

        if "quality_dimensions" in config:
            system_prompt += "\nQuality Dimensions:\n"
            for key, value in config["quality_dimensions"].items():
                system_prompt += f"- {key}: {value}\n"

        # 添加反幻觉约束
        if "anti_hallucination_constraints" in config:
            system_prompt += "\nAnti-Hallucination Constraints:\n"
            constraints = config["anti_hallucination_constraints"]
            for key, rules in constraints.items():
                system_prompt += f"\n{key}:\n"
                if isinstance(rules, list):
                    for rule in rules:
                        system_prompt += f"  - {rule}\n"

        # 添加方法论
        if "diagnostic_methodology" in config:
            system_prompt += "\nMethodology:\n"
            system_prompt += yaml.dump(config["diagnostic_methodology"], allow_unicode=True)

        return system_prompt.strip()

    def get_prompt_version(self, prompt_id: str) -> str:
        """获取 Prompt 版本"""
        if prompt_id in self.index.get("prompts", {}):
            return self.index["prompts"][prompt_id].get("version", "unknown")
        return "unknown"

    def validate_prompt_file(self, prompt_id: str) -> bool:
        """
        验证 Prompt 文件是否存在且有效

        Args:
            prompt_id: Prompt 标识符

        Returns:
            True if valid, False otherwise
        """
        try:
            prompt_data = self.load_prompt(prompt_id)
            return bool(prompt_data["content"])
        except (ValueError, FileNotFoundError):
            return False


# 向后兼容的辅助函数
def load_prompt(prompt_id: str, prompts_dir: str = "prompts") -> str:
    """
    向后兼容的 Prompt 加载函数（返回 system prompt 字符串）

    Args:
        prompt_id: Prompt 标识符
        prompts_dir: Prompt 目录

    Returns:
        构建的 system prompt 字符串
    """
    loader = PromptLoader(prompts_dir)

    try:
        return loader.build_system_prompt(prompt_id)
    except (ValueError, FileNotFoundError) as e:
        print(f"[PromptLoader] 无法加载 Prompt '{prompt_id}': {e}")
        print(f"[PromptLoader] 使用默认 Prompt")
        return f"You are an AI assistant specialized in {prompt_id}."


if __name__ == "__main__":
    # 测试加载器
    print("=== Prompt Loader Test ===\n")

    loader = PromptLoader()

    # 列出所有 Prompt
    print("Available Prompts:")
    for prompt in loader.list_prompts():
        print(f"  - {prompt['id']}: {prompt['description']}")

    print("\nAvailable Workflows:")
    for workflow in loader.list_workflows():
        print(f"  - {workflow['id']}: {workflow['name']} ({workflow['steps']} steps)")

    # 测试加载 Prompt
    print("\n=== Loading Prompt: ai_auto_self_diagnose ===")
    try:
        prompt_data = loader.load_prompt("ai_auto_self_diagnose")
        print(f"Role: {prompt_data['role']}")
        print(f"Version: {prompt_data['version']}")
        print(f"File: {prompt_data['file']}")
        print(f"Checksum: {prompt_data['checksum']}")
        print(f"Content length: {len(prompt_data['content'])} characters")

        # 构建 system prompt
        system_prompt = loader.build_system_prompt("ai_auto_self_diagnose")
        print(f"\nSystem Prompt (first 200 chars):\n{system_prompt[:200]}...")

    except Exception as e:
        print(f"Error: {e}")
